package cn.corffen.test.algorithm.st;

public class RBT<Key extends Comparable<Key>, Value> {

    public static final boolean RED = true;
    public static final boolean BLACK = false;

    private Node root;

    private class Node {
        private Node left;
        private Node right;
        private int size;
        private boolean color;
        private Key key;
        private Value val;

        public Node(int size, boolean color, Key key, Value val) {
            this.size = size;
            this.color = color;
            this.key = key;
            this.val = val;
        }
    }

    public int size() {
        return size(root);
    }

    public int size(Node x) {
        if (x == null) {
            return 0;
        }
        return x.size;
    }

    public boolean isEmpty() {
        return size() == 0;
    }

    public boolean contains(Key key) {
        if (key == null) {
            return false;
        }
        return get(key) != null;
    }

    public Value get(Key key) {
        if (key == null) {
            throw new IllegalArgumentException("key is null");
        }
        Node x = root;
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp < 0) {
                x = x.left;
            } else if (cmp > 0) {
                x = x.right;
            } else {
                return x.val;
            }
        }
        return null;
    }

    public void put(Key key, Value value) {
        if (key == null) {
            throw new IllegalArgumentException("key with put() is null");
        }
        root = put(key, value, root);
        root.color = BLACK;
    }

    public Node put(Key key, Value val, Node x) {
        if (x == null) {
            return new Node(1, RED, key, val);
        }
        int cmp = key.compareTo(x.key);
        if (cmp > 0) {
            put(key, val, x.right);
        } else if (cmp < 0) {
            put(key, val, x.left);
        } else {
            x.val = val;
        }

        if (isRed(x.right) && !isRed(x.left)) {
            x = rotateLeft(x);
        }
        if (isRed(x.left) && isRed(x.left.left)) {
            x = rotateRight(x);
        }
        if (isRed(x.left) && isRed(x.right)) {
            flipColors(x);
        }
        return x;
    }

    public void flipColors(Node h) {
        h.color = !h.color;
        h.left.color = !h.left.color;
        h.right.color = !h.right.color;
    }

    public Node rotateLeft(Node h) {
        Node x = h.right;
        x.left = h;
        h.right = x.left;
        x.color = h.color;
        h.color = RED;
        x.size = h.size;
        h.size = 1 + size(h.left) + size(h.right);
        return x;
    }

    public Node rotateRight(Node h) {
        Node x = h.left;
        x.right = h;
        h.left = x.right;
        x.size = h.size;
        h.size = 1 + size(h.left) + size(h.right);
        x.color = h.color;
        h.color = RED;
        return x;
    }

    public boolean isRed(Node x) {
        if (x == null) {
            return false;
        }
        return x.color;
    }
}
